// 研究了一天的状态机，天气也热，正当我写的即将昏迷之际，我突然醒悟了，原来状态机是这么一回事
// - 状态比喻成 数据结构
// - 事件比喻成 用户输入
// - 状态转移则是函数调用
// 如此依赖写成函数，也就是 (Orz 原来如此)
// type FSM struct {
// 	State 			FSMState
// 	TransformFuncs  map[FSMState] func()
// }

// func (f *FSM) UserAction(action FSMAction) {
// 	...
// }

package main

import (
	"errors"
	"fmt"
	"github.com/franela/goreq"
	"github.com/shirou/gopsutil/process"
	"io"
	"io/ioutil"
	"os"
	"os/user"
	"path/filepath"
	"runtime"
	"strings"
	"sync"
	"syscall"
	"time"

	"github.com/codeskyblue/kexec"
	"github.com/kennygrant/sanitize"
	log "github.com/linxlib/logs"
)

type FSMState string
type FSMEvent string
type FSMHandler func()

type FSM struct {
	mu       sync.Mutex
	state    FSMState
	handlers map[FSMState]map[FSMEvent]FSMHandler

	StateChange func(oldState, newState FSMState)
}

func (f *FSM) AddHandler(state FSMState, event FSMEvent, hdlr FSMHandler) *FSM {
	_, ok := f.handlers[state]
	if !ok {
		f.handlers[state] = make(map[FSMEvent]FSMHandler)
	}
	if _, ok = f.handlers[state][event]; ok {
		log.Fatalf("set twice for state(%s) event(%s)", state, event)
	}
	f.handlers[state][event] = hdlr
	return f
}

func (f *FSM) State() FSMState {
	return f.state
}

func (f *FSM) SetState(newState FSMState) {
	if f.StateChange != nil {
		f.StateChange(f.state, newState)
	}
	f.state = newState
}

func (f *FSM) Operate(event FSMEvent) FSMState {
	f.mu.Lock()
	defer f.mu.Unlock()

	eventMap := f.handlers[f.State()]
	if eventMap == nil {
		return f.State()
	}
	if fn, ok := eventMap[event]; ok {
		fn()
	}
	return f.State()
}

func NewFSM(initState FSMState) *FSM {
	return &FSM{
		state:    initState,
		handlers: make(map[FSMState]map[FSMEvent]FSMHandler),
	}
}

// Only 4 states now is enough, I think
// 2016-09-18 now five
var (
	Running   = FSMState("运行中")
	Stopped   = FSMState("已停止")
	Fatal     = FSMState("错误")
	RetryWait = FSMState("重试等待")
	Stopping  = FSMState("停止中")

	StartEvent   = FSMEvent("启动")
	StopEvent    = FSMEvent("停止")
	RestartEvent = FSMEvent("重启")
)

type Program struct {
	Name          string   `yaml:"name" json:"name"`
	Command       string   `yaml:"command" json:"command"`
	Environ       []string `yaml:"environ" json:"environ"`
	Dir           string   `yaml:"directory" json:"directory"`
	StartAuto     bool     `yaml:"start_auto" json:"startAuto"`
	StartRetries  int      `yaml:"start_retries" json:"startRetries"`
	StartSeconds  int      `yaml:"start_seconds,omitempty" json:"startSeconds"`
	StopTimeout   int      `yaml:"stop_timeout,omitempty" json:"stopTimeout"`
	User          string   `yaml:"user,omitempty" json:"user"`
	Notifications struct {
		Bark Bark  `yaml:"bark,omitempty" json:"bark,omitempty"`
	} `yaml:"notifications,omitempty" json:"-"`
}

type Bark struct {
	Enable bool `yaml:"enable" json:"enable"`
	Users  []string `yaml:"users" json:"users"`
}

func (p *Program) Check() error {
	if p.Name == "" {
		return errors.New("Program name empty")
	}
	if p.Command == "" {
		return errors.New("Program command empty")
	}
	return nil
}

func (p *Program) IsRunning() (int,bool)  {
  	pss,err:= process.Processes()
	if err != nil {
		return -1,false
	}
	for _,v:=range pss {
		cmdline,_ := v.Name()
		if strings.Contains(p.Command,cmdline) {
			return int(v.Pid),true
		}
	}
	return -1,false
}

func (p *Program) RunNotification() {
	if p.Notifications.Bark.Enable {
		for _,key := range p.Notifications.Bark.Users {
			msg:=fmt.Sprintf("%s 意外退出", p.Name)
			getUrl := fmt.Sprintf("https://api.day.app/%s/%s/%s",key,msg,time.Now().Format(time.RFC3339))
			goreq.Request{ Uri: getUrl }.Do()
		}
	}
}

func IsRoot() bool {
	u, err := user.Current()
	return err == nil && u.Username == "root"
}

type Process struct {
	*FSM       `json:"-"`
	Program    `json:"program"`
	cmd        *kexec.KCommand
	Stdout     *QuickLossBroadcastWriter `json:"-"`
	Stderr     *QuickLossBroadcastWriter `json:"-"`
	Output     *QuickLossBroadcastWriter `json:"-"`
	OutputFile *os.File                  `json:"-"`
	stopC      chan syscall.Signal
	retryLeft  int
	Status     string `json:"status"`

	mu sync.Mutex
}

// FIXME(ssx): maybe need to return error
func (p *Process) buildCommand() *kexec.KCommand {
	cmd := kexec.CommandString(p.Command)
	// cmd := kexec.Command(p.Command[0], p.Command[1:]...)
	logDir := filepath.Join(defaultConfigDir, "log", sanitize.Name(p.Name))
	if !IsDir(logDir) {
		os.MkdirAll(logDir, 0755)
	}
	var fout io.Writer
	var err error
	p.OutputFile, err = os.OpenFile(filepath.Join(logDir, "output.log"), os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0644)
	if err != nil {
		log.Warn("创建 stdout 文件失败:", err)
		fout = ioutil.Discard
	} else {
		fout = p.OutputFile
	}
	cmd.Stdout = io.MultiWriter(p.Stdout, p.Output, fout)
	cmd.Stderr = io.MultiWriter(p.Stderr, p.Output, fout)
	// config environ
	cmd.Env = os.Environ() // inherit current vars
	environ := map[string]string{}
	if p.User != "" {
		if !IsRoot() {
			log.Warnf("检测到当前不是root用户, 无法切换用户")
		} else if err := cmd.SetUser(p.User); err != nil {
			log.Warnf("[%s] 切换用户到 %s 失败, %v", p.Name, p.User, err)
		} else {
			var homeDir string
			switch runtime.GOOS {
			case "linux":
				homeDir = "/home/" + p.User // FIXME(ssx): maybe there is a better way
			case "darwin":
				homeDir = "/Users/" + p.User
			}
			cmd.Env = append(cmd.Env, "HOME="+homeDir, "USER="+p.User)
			environ["HOME"] = homeDir
			environ["USER"] = p.User
		}
	}
	cmd.Env = append(cmd.Env, p.Environ...)
	mapping := func(key string) string {
		val := os.Getenv(key)
		if val != "" {
			return val
		}
		return environ[key]
	}
	cmd.Dir = os.Expand(p.Dir, mapping)
	if strings.HasPrefix(cmd.Dir, "~") {
		cmd.Dir = mapping("HOME") + cmd.Dir[1:]
	}
	log.Infof("[%s] 使用目录: %s\n", p.Name, cmd.Dir)
	return cmd
}

func (p *Process) waitNextRetry() {
	p.SetState(RetryWait)
	if p.retryLeft <= 0 {
		p.retryLeft = p.StartRetries
		p.SetState(Fatal)
		return
	}
	p.retryLeft -= 1
	select {
	case <-time.After(2 * time.Second): // TODO: need put it into Program
		p.startCommand()
	case <-p.stopC:
		p.stopCommand()
	}
}

func (p *Process) stopCommand() {
	p.mu.Lock()
	defer p.mu.Unlock()
	defer p.SetState(Stopped)
	if p.cmd == nil {
		return
	}
	p.SetState(Stopping)
	if p.cmd.Process != nil {
		p.cmd.Process.Signal(syscall.SIGTERM) // TODO(ssx): add it to config
	}
	select {
	case <-GoFunc(p.cmd.Wait):
		log.Printf("程序(%s) 正常退出", p.Name)
	case <-time.After(time.Duration(p.StopTimeout) * time.Second): // TODO: add 3s to config
		log.Printf("程序(%s) 关闭全部", p.Name)
		p.cmd.Terminate(syscall.SIGKILL) // cleanup
	}
	err := p.cmd.Wait() // This is OK, because Signal KILL will definitely work
	prefixStr := "\n--- GOSUV LOG " + time.Now().Format("2006-01-02 15:04:05")
	if err == nil {
		io.WriteString(p.cmd.Stderr, fmt.Sprintf("%s 退出成功 ---\n\n", prefixStr))
	} else {
		io.WriteString(p.cmd.Stderr, fmt.Sprintf("%s 退出 %v ---\n\n", prefixStr, err))
	}
	if p.OutputFile != nil {
		p.OutputFile.Close()
		p.OutputFile = nil
	}
	p.cmd = nil
}

func (p *Process) IsRunning() bool {
	return p.State() == Running || p.State() == RetryWait
}

func (p *Process) startCommand() {
	log.Printf("运行命令(%s): %s", p.Name, p.Command)
	p.cmd = p.buildCommand()

	p.SetState(Running)

	if pid,running:=p.Program.IsRunning();running {
		log.Infof("程序 %s 已运行, 结束进程 %d", p.Name,pid)
		pro,_:=os.FindProcess(pid)
		err := pro.Kill()
		pro.Release()
		if err != nil {
			log.Error(err)
		}
	}
	if err := p.cmd.Start(); err != nil {
		log.Warnf("程序 %s 启动失败: %v", p.Name, err)
		p.SetState(Fatal)
		return
	}

	go func() {
		errC := GoFunc(p.cmd.Wait)
		startTime := time.Now()
		select {
		case <-errC:
			//由于程序自身或外部原因退出
			go p.RunNotification()
			// if p.cmd.Wait() returns, it means program and its sub process all quited. no need to kill again
			// func Wait() will only return when program session finishs. (Only Tested on mac)
			log.Printf("程序(%s) 结束, 耗时 %v", p.Name, time.Since(startTime))
			if time.Since(startTime) < time.Duration(p.StartSeconds)*time.Second {
				if p.retryLeft == p.StartRetries { // If first time quit so fast, just set to fatal
					p.SetState(Fatal)
					log.Printf("程序(%s) 启动后很快退出, 状态 -> 错误", p.Name)
					return
				}
			}
			p.waitNextRetry()
		case <-p.stopC:
			log.Println("接收停止命令")
			p.stopCommand() // clean up all process
		}
	}()
}

func NewProcess(pg Program) *Process {
	outputBufferSize := 24 * 1024 // 24K
	pr := &Process{
		FSM:       NewFSM(Stopped),
		Program:   pg,
		stopC:     make(chan syscall.Signal),
		retryLeft: pg.StartRetries,
		Status:    string(Stopped),
		Output:    NewQuickLossBroadcastWriter(outputBufferSize),
		Stdout:    NewQuickLossBroadcastWriter(outputBufferSize),
		Stderr:    NewQuickLossBroadcastWriter(outputBufferSize),
	}

	pr.StateChange = func(_, newStatus FSMState) {
		pr.Status = string(newStatus)

		// TODO: status need to filter with config, not hard coded.
		//if newStatus == Fatal {
		//	go pr.Program.RunNotification(errors.New("newStatus is fatal"))
		//}
	}
	if pr.StartSeconds <= 0 {
		pr.StartSeconds = 3
	}
	if pr.StopTimeout <= 0 {
		pr.StopTimeout = 3
	}

	pr.AddHandler(Stopped, StartEvent, func() {
		pr.retryLeft = pr.StartRetries
		pr.startCommand()
	})
	pr.AddHandler(Fatal, StartEvent, pr.startCommand)

	pr.AddHandler(Running, StopEvent, func() {
		select {
		case pr.stopC <- syscall.SIGTERM:
		case <-time.After(200 * time.Millisecond):
		}
	}).AddHandler(Running, RestartEvent, func() {
		go func() {
			pr.Operate(StopEvent)
			// TODO: start laterly
			time.Sleep(1 * time.Second)
			pr.Operate(StartEvent)
		}()
	})
	return pr
}


